package org.micro.neural.config.store;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

import io.lettuce.core.RedisClient;
import io.lettuce.core.RedisURI;
import io.lettuce.core.api.StatefulRedisConnection;
import io.lettuce.core.api.sync.RedisCommands;
import io.lettuce.core.pubsub.StatefulRedisPubSubConnection;
import io.lettuce.core.pubsub.api.async.RedisPubSubAsyncCommands;
import io.lettuce.core.support.ConnectionPoolSupport;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.pool2.impl.GenericObjectPool;
import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
import org.micro.neural.URL;
import org.micro.neural.common.Constants;
import org.micro.neural.extension.Extension;

/**
 * The Store by Redis
 *
 * @author lry
 **/
@Slf4j
@Extension("redis")
public class RedisStore implements IStore {

    private RedisClient redisClient = null;
    private GenericObjectPool<StatefulRedisConnection<String, String>> genericObjectPool;
    private final Map<IStoreListener, RedisPubSubAsyncCommands<String, String>> subscribed = new ConcurrentHashMap<>();

    @Override
    public void initialize(URL url) {
        RedisURI redisURI = RedisURI.Builder.redis(url.getHost(), url.getPort()).build();
        this.redisClient = RedisClient.create(redisURI);
        this.genericObjectPool = ConnectionPoolSupport.createGenericObjectPool(
                () -> redisClient.connect(), new GenericObjectPoolConfig());
    }

    @Override
    public void batchUpOrAdd(long expire, Map<String, Long> data) {
        try {
            try (StatefulRedisConnection<String, String> connection = genericObjectPool.borrowObject()) {
                RedisCommands<String, String> commands = connection.sync();
                for (Map.Entry<String, Long> entry : data.entrySet()) {
                    commands.incrby(entry.getKey(), entry.getValue());
                    commands.pexpire(entry.getKey(), expire);
                }
            }
        } catch (Exception e) {
            log.error(e.getMessage(), e);
        }
    }

    @Override
    public void add(ModuleKey moduleKey, StoreKey storeKey, Object data) {
        try {
            try (StatefulRedisConnection<String, String> connection = genericObjectPool.borrowObject()) {
                RedisCommands<String, String> commands = connection.sync();
                commands.hset(moduleKey.serialize(), storeKey.serialize(), Constants.serialize(data));
            }
        } catch (Exception e) {
            log.error(e.getMessage(), e);
        }
    }

    @Override
    public Set<StoreKey> searchKeys(ModuleKey moduleKey, String keyword) {
        Set<StoreKey> storeKeySet = new HashSet<>();
        try {
            try (StatefulRedisConnection<String, String> connection = genericObjectPool.borrowObject()) {
                RedisCommands<String, String> commands = connection.sync();
                List<String> keys = commands.hkeys(moduleKey.serialize());
                if (keys == null || keys.isEmpty()) {
                    return storeKeySet;
                }

                for (String key : keys) {
                    storeKeySet.add(new StoreKey().deserialize(key));
                }

                return storeKeySet;
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public <C> C query(ModuleKey moduleKey, StoreKey storeKey, Class<C> clz) {
        try {
            try (StatefulRedisConnection<String, String> connection = genericObjectPool.borrowObject()) {
                RedisCommands<String, String> commands = connection.sync();
                return Constants.deserialize(clz, commands.hget(moduleKey.serialize(), storeKey.serialize()));
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public void publish(ModuleKey moduleKey, StoreKey storeKey, Object data) {
        String channel = moduleKey.serialize() + Constants.DELIMITER + storeKey.serialize();
        try {
            try (StatefulRedisConnection<String, String> connection = genericObjectPool.borrowObject()) {
                RedisCommands<String, String> commands = connection.sync();
                commands.publish(channel, Constants.serialize(data));
            }
        } catch (Exception e) {
            log.error(e.getMessage(), e);
        }
    }

    @Override
    public void subscribe(ModuleKey moduleKey, StoreKey[] storeKeys, IStoreListener listener) {
        List<String> channels = new ArrayList<>();
        for (StoreKey storeKey : storeKeys) {
            channels.add(moduleKey.serialize() + Constants.DELIMITER + storeKey.serialize());
        }

        StatefulRedisPubSubConnection<String, String> connection = redisClient.connectPubSub();
        connection.addListener(null);
        RedisPubSubAsyncCommands<String, String> commands = connection.async();

        subscribed.put(listener, commands);
        commands.subscribe(channels.toArray(new String[0]));
    }

    @Override
    public void unSubscribe(IStoreListener listener) {
        RedisPubSubAsyncCommands<String, String> commands = subscribed.get(listener);
        if (commands != null) {
            commands.unsubscribe();
            subscribed.remove(listener);
        }
    }

    @Override
    public void destroy() {
        if (null != genericObjectPool) {
            genericObjectPool.close();
        }
        if (null != redisClient) {
            redisClient.shutdown();
        }
    }

}
